from paddleocr import PPDocTranslation
import os
import time
import sys
import logging
import dotenv
from pathlib import Path
from typing import Dict, List, Optional, Union
import shutil

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 加载环境变量
dotenv.load_dotenv()

def get_chat_bot_config() -> Dict:
    """获取大模型配置"""
    return {
        "module_name": os.getenv("MODULE_NAME", "Qwen/Qwen2.5-VL-72B-Instruct-AWQ"),
        "model_name": os.getenv("MODEL_NAME", "Qwen/Qwen2.5-VL-72B-Instruct-AWQ"),
        "base_url": os.getenv("BASE_URL", ""),
        "api_type": os.getenv("API_TYPE", "openai"),
        "api_key": os.getenv("API_KEY", "dummy-key"),
    }

def extract_text_from_md_info(md_info):
    """尝试从markdown信息对象中提取文本内容，优先获取markdown_texts属性"""
    # 优先获取markdown_texts属性
    if hasattr(md_info, 'markdown_texts') and isinstance(md_info.markdown_texts, str):
        return md_info.markdown_texts
    if hasattr(md_info, 'markdown_text') and isinstance(md_info.markdown_text, str):
        return md_info.markdown_text
    if hasattr(md_info, 'text') and isinstance(md_info.text, str):
        return md_info.text
    if hasattr(md_info, 'content') and isinstance(md_info.content, str):
        return md_info.content
    if hasattr(md_info, 'markdown') and isinstance(md_info.markdown, str):
        return md_info.markdown
    
    # 如果无法通过属性获取，则尝试将对象转换为字符串
    try:
        text_content = str(md_info)
        if text_content and len(text_content) > 10:  # 确保内容有意义
            return text_content
    except:
        pass
    
    return "# 翻译结果\n\n无法提取原文内容"

def save_md_result_to_file(md_info, output_file_path):
    """将markdown结果保存到文件，优先处理markdown_texts属性"""
    try:
        # 尝试使用save_to_markdown方法前，先尝试提取markdown_texts
        markdown_content = None
        
        # 优先尝试获取markdown_texts内容
        if hasattr(md_info, 'markdown_texts') and isinstance(md_info.markdown_texts, str):
            markdown_content = md_info.markdown_texts
        elif hasattr(md_info, 'markdown_text') and isinstance(md_info.markdown_text, str):
            markdown_content = md_info.markdown_text
            
        # 如果成功获取了markdown_texts内容，直接保存
        if markdown_content:
            with open(output_file_path, 'w', encoding='utf-8') as f:
                f.write(markdown_content)
            logger.info(f"直接保存markdown_texts到文件: {output_file_path}")
            return output_file_path
        
        # 如果没有markdown_texts，尝试使用对象自己的save_to_markdown方法
        if hasattr(md_info, 'save_to_markdown') and callable(md_info.save_to_markdown):
            try:
                saved_path = md_info.save_to_markdown(os.path.dirname(output_file_path))
                if saved_path and os.path.exists(saved_path):
                    return saved_path
            except Exception as e:
                logger.warning(f"使用save_to_markdown方法保存失败: {e}")
        
        # 如果上面的方法都失败，提取文本并手动保存
        text_content = extract_text_from_md_info(md_info)
        
        # 确保文本是字符串
        if not isinstance(text_content, str):
            text_content = f"# 内容类型错误\n\n类型: {type(text_content)}"
        
        with open(output_file_path, 'w', encoding='utf-8') as f:
            f.write(text_content)
        
        return output_file_path
    except Exception as e:
        logger.error(f"保存文件失败: {e}")
        return None

def translate_document(
    input_path: str, 
    output_path: str = "./output", 
    target_language: str = "zh",
    chunk_size: int = 5000
) -> Dict:
    """
    翻译文档（支持PDF和Markdown文件）
    
    Args:
        input_path: 输入文件路径
        output_path: 输出目录
        target_language: 目标语言（如：zh, en等）
        chunk_size: 翻译块大小
        
    Returns:
        Dict: 包含处理结果的字典
    """
    results = {
        "success": False,
        "message": "",
        "output_files": [],
        "file_urls": [],
        "error": None
    }
    
    try:
        # 确保输出目录存在
        os.makedirs(output_path, exist_ok=True)
        
        logger.info(f"初始化PPDocTranslation...")
        pipeline = PPDocTranslation()
        logger.info(f"PPDocTranslation初始化完成！")
        
        # 获取聊天机器人配置
        chat_bot_config = get_chat_bot_config()
        logger.info(f"使用API: {chat_bot_config['base_url']}")
        
        logger.info(f"开始处理文件：{input_path}")
        
        if input_path.lower().endswith(".md"):
            # 处理Markdown文件
            logger.info("开始加载Markdown文件...")
            ori_md_info_list = pipeline.load_from_markdown(input_path)
            logger.info("Markdown文件加载完成！")
        else:
            # 处理PDF/图像文件
            logger.info("开始使用PP-StructureV3进行布局分析...")
            visual_predict_res = pipeline.visual_predict(
                input_path,
                use_doc_orientation_classify=False,
                use_doc_unwarping=False,
                use_common_ocr=True,
                use_seal_recognition=True,
                use_table_recognition=True,
            )
            logger.info("PP-StructureV3布局分析完成！")
            
            ori_md_info_list = []
            for i, res in enumerate(visual_predict_res):
                logger.info(f"处理第 {i+1}/{len(visual_predict_res)} 页...")
                layout_parsing_result = res["layout_parsing_result"]
                ori_md_info_list.append(layout_parsing_result.markdown)
                
                # 保存图片和Markdown
                layout_parsing_result.save_to_img(output_path)
                layout_parsing_result.save_to_markdown(output_path)
                logger.info(f"第 {i+1} 页处理完成")
            
            # 合并多页文档
            if input_path.lower().endswith(".pdf"):
                logger.info("合并多页文档的Markdown信息...")
                ori_md_info = pipeline.concatenate_markdown_pages(ori_md_info_list)
                logger.info("保存合并后的原始Markdown文本...")
                ori_md_info.save_to_markdown(output_path)
                logger.info("保存完成！")
        
        # 保存原始文本，以便在翻译失败时使用
        original_texts = []
        original_files = []
        for i, md_info in enumerate(ori_md_info_list):
            # 优先获取markdown_texts
            text = extract_text_from_md_info(md_info)
            original_texts.append(text)
            
            # 保存原始文本文件
            base_name = os.path.basename(input_path)
            name_without_ext = os.path.splitext(base_name)[0]
            original_filename = f"{name_without_ext}_original_{i+1}.md"
            original_filepath = os.path.join(output_path, original_filename)
            
            try:
                with open(original_filepath, 'w', encoding='utf-8') as f:
                    f.write(text)
                original_files.append(original_filepath)
            except Exception as e:
                logger.error(f"保存原始文本时出错: {str(e)}")
        
        # 进行翻译
        logger.info("开始进行文档翻译...")
        logger.info(f"目标语言：{target_language}")
        try:
            tgt_md_info_list = pipeline.translate(
                ori_md_info_list=ori_md_info_list,
                target_language=target_language,
                chunk_size=chunk_size,
                chat_bot_config=chat_bot_config,
            )
        except Exception as e:
            logger.error(f"翻译过程出错: {str(e)}")
            # 如果翻译失败，我们仍然保存原始文本
            tgt_md_info_list = ori_md_info_list
            results["error"] = f"翻译失败，将返回原始文本: {str(e)}"
        
        # 保存翻译结果
        logger.info("准备保存翻译或原始文本...")
        for i, (tgt_md_info, original_text) in enumerate(zip(tgt_md_info_list, original_texts)):
            logger.info(f"保存第 {i+1}/{len(tgt_md_info_list)} 个文件...")
            
            # 创建文件名
            base_name = os.path.basename(input_path)
            name_without_ext = os.path.splitext(base_name)[0]
            result_filename = f"{name_without_ext}_{target_language}_{i+1}.md"
            output_file_path = os.path.join(output_path, result_filename)
            
            try:
                # 尝试检查并打印tgt_md_info的属性，用于调试
                logger.info(f"翻译结果对象类型: {type(tgt_md_info)}")
                if hasattr(tgt_md_info, 'markdown_texts'):
                    logger.info(f"发现markdown_texts属性")
                if hasattr(tgt_md_info, 'markdown_text'):
                    logger.info(f"发现markdown_text属性")
                if hasattr(tgt_md_info, 'text'):
                    logger.info(f"发现text属性")
                
                # 使用我们的辅助函数保存文件
                saved_path = save_md_result_to_file(tgt_md_info, output_file_path)
                
                if saved_path and os.path.exists(saved_path):
                    results["output_files"].append(str(saved_path))
                    logger.info(f"成功保存翻译结果到: {saved_path}")
                else:
                    # 如果保存失败，使用原始文本
                    with open(output_file_path, 'w', encoding='utf-8') as f:
                        f.write(original_text)
                    results["output_files"].append(output_file_path)
                    logger.info(f"使用原始文本作为备用: {output_file_path}")
            except Exception as e:
                logger.error(f"保存翻译结果时出错: {str(e)}")
                # 添加原始文件作为备用
                if i < len(original_files) and os.path.exists(original_files[i]):
                    results["output_files"].append(original_files[i])
                    logger.info(f"添加原始文件作为备用: {original_files[i]}")
        
        logger.info("翻译结果保存完成！")
        results["success"] = True
        results["message"] = "文档处理和翻译完成"
        
    except Exception as e:
        error_msg = f"处理过程中出错：{str(e)}"
        logger.error(error_msg)
        import traceback
        logger.error(traceback.format_exc())
        results["success"] = False
        results["message"] = error_msg
        results["error"] = str(e)
    
    return results

# 测试代码
if __name__ == "__main__":
    # 测试PDF文件翻译
    result = translate_document("input_files/DeepSeek-V3-2pages.pdf", output_path="output_files/temp/test")
    print(result) 